#include <cassert>
#include <cstdlib>
#include <cmath>
#include <ctime>
#include <vector>
#include <iostream>

#include <limits>

#include "sds/sfile.h"
#include "sds/sdataview.h"

#include "cluster.h"

#include "random.h"

using namespace std;
using namespace sds;

typedef SDataView<unsigned char> CharView;
typedef SDataView<Point> PointView;

const unsigned int state = 50;
const double maxdist = 80.0;

Random rng;

void Read(const string& filename, PointView& view, vector<SInteger>& size) {
  SFile file(filename);

  CharView source[CH];

  for (int k = 0; k < CH; ++k) {
    source[k] = file.getDataset(k).getDataView<unsigned char>();
  }
  
  size = source[0].getSize();
  
  view = PointView(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      view[i][j] = point;
    }
  }
}



int main(int argc, const char * argv[]) {
  assert(argc == 3);

  PointView image;
  vector<SInteger> size;

  Read(argv[1], image, size);

  CharView point(size);

  Cluster cluster[state];

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      char cl = (char)(state * rng.uniform()); 
      point[i][j] = cl + 1;
      cluster[(int)cl] += image[i][j];
    }
  }
  int x = (int)(rng.uniform() * size[0]);
  int y = (int)(rng.uniform() * size[1]);
  const int length = 60;
  for (int k = 0; k < length; ++k) {
    cout << "Step: " << k << endl;
    double t = 3.0 / log (1.0 + k * 0.08);
    for (int i = 0; i < size[0]; ++i) {
      for (int j = 0; j < size[1]; ++j) {
	if (point[i][j] != 0) {
	  cluster[point[i][j] - 1] -= image[i][j];
	}
	double probs[state];
	double sum = 0.0;
	for (int l = 0; l < (int)state; ++l) {
	  double dist = distance_rp(cluster[l].center(), image[i][j]);
	  probs[l] = 
	    exp( - dist / t);
	  sum += probs[l];
	}
	if (i == x && j == y) {
	  for (int l = 0; l < (int)state; ++l) {
	    cout << probs[l] << ' ';
	  }
	  cout << exp( - maxdist / t) << endl;
	}
	sum += exp( - maxdist / t);
	double prob = rng.uniform() * sum;
	point[i][j] = 0;
	for (int l = 0; l < (int)state; ++l) {
	  if (prob < probs[l]) {
	    point[i][j] = l + 1;
	    cluster[l] += image[i][j];
	    break;
	  }
	  prob -= probs[l];
	}
      }
    }
  } 

  SFile output(argv[2], SFile::CREATE);
  SDataset out = output.createDataset("Clustered Image", UINT8, size);
  out.setDataView<unsigned char>(point);
  return 0;
}
